#pragma once
#include "file.h"
#include "task.h"
#include "uring_loop.h"
#include <arpa/inet.h>
#include <asm-generic/socket.h>
#include <cstdint>
#include <cstring>
#include <linux/close_range.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <variant>

namespace co_async {

class SocketHandle : public FileHandle {
public:
  explicit SocketHandle(int fd = -1) : FileHandle(fd) {}
  SocketHandle(SocketHandle &&that) : FileHandle(std::move(that)) {}
  SocketHandle &operator=(SocketHandle &&that) {
    FileHandle::operator=(std::move(that));
    return *this;
  }

  static Task<SocketHandle> CreateSocket(int family, int type) {
    int fd = co_await uring_socket(loop, family, type, 0, 0);
    // std::cout << "create sock" << std::endl;
    co_return SocketHandle{fd};
  }
};

class IpAddress {
public:
  friend class SocketAddress;
  explicit IpAddress(struct in_addr const &addr) noexcept : m_addr(addr) {}
  explicit IpAddress(struct in6_addr const &addr6) noexcept : m_addr(addr6) {}
  IpAddress(const char *ip) {
    in_addr addr;
    if (check_error(inet_pton(AF_INET, ip, &addr))) {
      m_addr = addr;
      return;
    }
    in6_addr addr6;
    if (check_error(inet_pton(AF_INET6, ip, &addr6))) {
      m_addr = addr6;
      return;
    }
    auto hent = gethostbyname(ip);
    for (int i = 0; hent->h_addr_list[i]; i++) {
      if (hent->h_addrtype == AF_INET) {
        std::memcpy(&addr, hent->h_addr_list[i], sizeof(in_addr));
        m_addr = addr;
        return;
      } else if (hent->h_addrtype == AF_INET6) {
        std::memcpy(&addr6, hent->h_addr_list[i], sizeof(in6_addr));
        m_addr = addr6;
        return;
      }
    }
    throw std::invalid_argument("invalid domain name or ip address");
  }

private:
  std::variant<in_addr, in6_addr> m_addr;
};

class SocketAddress {
public:
  friend class SocketServer;

  SocketAddress(const SocketAddress &addr) = default;
  SocketAddress(IpAddress ip, uint16_t port) : m_port(port), m_ip(ip) {
    std::visit([&](auto const &addr) { init_addr(addr, port); }, ip.m_addr);
  }
  sa_family_t get_family() const { return m_family; };

  socklen_t get_addr_len() const { return m_addr_len; }
  uint16_t get_port() { return m_port; }
  IpAddress get_ip() const { return m_ip; }

private:
  void init_addr(const in_addr &addr, uint16_t port) {
    sockaddr_in saddr;
    m_family = saddr.sin_family = AF_INET;
    saddr.sin_port = htons(port);
    std::memcpy(&saddr.sin_addr, &addr, sizeof(in_addr));
    m_addr = saddr;
    m_addr_len = sizeof(saddr);
  }
  void init_addr(const in6_addr &addr, uint16_t port) {
    sockaddr_in6 saddr;
    m_family = saddr.sin6_family = AF_INET6;
    saddr.sin6_port = htons(port);
    std::memcpy(&saddr.sin6_addr, &addr, sizeof(in_addr));
    m_addr = saddr;
    m_addr_len = sizeof(saddr);
  }
  struct addr_visitor {
    sockaddr *operator()(sockaddr_in &addr) const {
      return reinterpret_cast<sockaddr *>(&addr);
    }
    sockaddr *operator()(sockaddr_in6 &addr) const {
      return reinterpret_cast<sockaddr *>(&addr);
    }
  };
  sockaddr *get_addr() { return std::visit(addr_visitor{}, m_addr); }
  sa_family_t m_family;
  std::variant<sockaddr_in, sockaddr_in6> m_addr;
  socklen_t m_addr_len;
  uint16_t m_port;
  IpAddress m_ip;
};

class SocketServer : SocketHandle {
public:
  SocketServer(const SocketAddress &addr, SocketHandle &&that)
      : m_addr(addr), SocketHandle(that.get_fd()) {
    that.release();
  };
  SocketServer(SocketServer &&that)
      : m_addr(that.m_addr), SocketHandle(std::move(that)) {}

  const SocketAddress &get_addr() const noexcept { return m_addr; }
  const sockaddr *get_sockaddr() { return m_addr.get_addr(); }
  socklen_t get_addr_len() const { return m_addr.get_addr_len(); }

  static Task<SocketServer> CreateServer(SocketAddress addr,
                                         int n = SOMAXCONN) {
    // std::cout << "sock" << std::endl;
    SocketHandle socket =
        co_await SocketHandle::CreateSocket(addr.get_family(), SOCK_STREAM);
    // int fd = socket(addr.get_family(), SOCK_STREAM, 0);
    int opt_val = 1;
    int fd = socket.get_fd();
    check_error(
        setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt_val, sizeof(opt_val)));
    check_error(bind(fd, addr.get_addr(), addr.get_addr_len()));
    check_error(listen(fd, n));
    co_return {addr, std::move(socket)};
  }

  static Task<SocketServer> CreateServer(IpAddress addr, uint16_t port,
                                         int n = SOMAXCONN) {
    SocketAddress saddr{addr, port};
    return CreateServer(saddr, n);
  }

  Task<SocketHandle> server_accept() {
    auto addr_len = m_addr.get_addr_len();
    int fd = co_await uring_accept(loop, m_fd, m_addr.get_addr(), &addr_len, 0);
    co_return SocketHandle{fd};
  }

private:
  SocketAddress m_addr;
};

} // namespace co_async